﻿#pragma once

#include "../../df.h"


namespace df
{
#include "aes.h"
#include "des.h"

	//算法标识
	struct CryptMode
	{
		enum{
			_Start = 1233,

			DES_CBC,
			Triple_DES_CBC,
			AES_CBC,
			AES256_CBC,

			_End,
		};
	};

	template<int Mode>
	struct CryptAlg
	{
		//static_assert(false, "no implement Crypt")
	};

	template<>
	struct CryptAlg<CryptMode::DES_CBC>
	{
		static const int KeySize = MBEDTLS_DES_KEY_SIZE;
		//初始化向量
		unsigned char iv_[KeySize];
		mbedtls_des_context contentEn_;
		mbedtls_des_context contentDe_;

		void InitByteKey(const unsigned char key[KeySize])
		{
			//初始化向量
			memcpy(iv_, key, KeySize);
			mbedtls_des_setkey_enc(&contentEn_, key);
			mbedtls_des_setkey_dec(&contentDe_, key);
		}

		void InitKey(const CC & password)
		{
			//初始化密钥
			unsigned char key[KeySize] = { 0 };
			CC::StringToByte(password, key, KeySize);

			//初始化向量
			memcpy(iv_, key, KeySize);

			mbedtls_des_setkey_enc(&contentEn_, key);
			mbedtls_des_setkey_dec(&contentDe_, key);
		}

		bool Encrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_des_crypt_cbc(&contentEn_, MBEDTLS_DES_ENCRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}

		bool Decrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_des_crypt_cbc(&contentDe_, MBEDTLS_DES_DECRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}
	};

	template<>
	struct CryptAlg < CryptMode::Triple_DES_CBC>
	{
		static const int KeySize = MBEDTLS_DES_KEY_SIZE * 2;
		//初始化向量
		unsigned char iv_[8];
		mbedtls_des3_context contentEn_;
		mbedtls_des3_context contentDe_;

		void InitByteKey(const unsigned char key[KeySize])
		{
			//初始化向量
			memcpy(iv_, key, KeySize);
			mbedtls_des3_set2key_enc(&contentEn_, key);
			mbedtls_des3_set2key_dec(&contentDe_, key);
		}

		void InitKey(const CC & password)
		{
			//初始化密钥
			unsigned char key[KeySize] = { 0 };
			CC::StringToByte(password, key, KeySize);

			//初始化向量
			memcpy(iv_, key, 8);

			mbedtls_des3_set2key_enc(&contentEn_, key);
			mbedtls_des3_set2key_dec(&contentDe_, key);
		}

		bool Encrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_des3_crypt_cbc(&contentEn_, MBEDTLS_DES_ENCRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}

		bool Decrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_des3_crypt_cbc(&contentDe_, MBEDTLS_DES_DECRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}
	};

	template<>
	struct CryptAlg < CryptMode::AES_CBC>
	{
		static const int KeySize = 16;
		//初始化向量
		unsigned char iv_[KeySize];
		mbedtls_aes_context contentEn_;
		mbedtls_aes_context contentDe_;

		void InitByteKey(const unsigned char key[KeySize])
		{
			//初始化向量
			memcpy(iv_, key, KeySize);
			mbedtls_aes_setkey_enc(&contentEn_, key, 128);
			mbedtls_aes_setkey_dec(&contentDe_, key, 128);
		}


		void InitKey(const CC & password)
		{
			//初始化密钥
			unsigned char key[KeySize] = { 0 };
			CC::StringToByte(password, key, KeySize);

			//初始化向量
			memcpy(iv_, key, KeySize);

			mbedtls_aes_setkey_enc(&contentEn_, key, 128);
			mbedtls_aes_setkey_dec(&contentDe_, key, 128);
		}

		bool Encrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_aes_crypt_cbc(&contentEn_, MBEDTLS_AES_ENCRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}

		bool Decrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_aes_crypt_cbc(&contentDe_, MBEDTLS_AES_DECRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}
	};

	template<>
	struct CryptAlg < CryptMode::AES256_CBC>
	{
		static const int KeySize = 32;
		//初始化向量
		unsigned char iv_[16];
		mbedtls_aes_context contentEn_;
		mbedtls_aes_context contentDe_;

		void InitByteKey(const unsigned char key[KeySize])
		{
			//初始化向量
			memcpy(iv_, key, 16);
			mbedtls_aes_setkey_enc(&contentEn_, key, 256);
			mbedtls_aes_setkey_dec(&contentDe_, key, 256);
		}


		void InitKey(const CC & password)
		{
			//初始化密钥
			unsigned char key[KeySize] = { 0 };
			CC::StringToByte(password, key, KeySize);

			//初始化向量
			memcpy(iv_, key, 16);

			mbedtls_aes_setkey_enc(&contentEn_, key, 256);
			mbedtls_aes_setkey_dec(&contentDe_, key, 256);
		}

		bool Encrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_aes_crypt_cbc(&contentEn_, MBEDTLS_AES_ENCRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}

		bool Decrypt(const void *input, void *output, size_t size)
		{
			return mbedtls_aes_crypt_cbc(&contentDe_, MBEDTLS_AES_DECRYPT, size, iv_, (unsigned char *)input, (unsigned char *)output) == 0;
		}
	};


	typedef CryptAlg<CryptMode::AES_CBC> CryptAES128;
	typedef CryptAlg<CryptMode::AES256_CBC> CryptAES256;
	typedef CryptAlg<CryptMode::DES_CBC> CryptDES;
	typedef CryptAlg<CryptMode::Triple_DES_CBC> CryptDES3;


	class Crypt
	{
	public:

		struct CryptHeader
		{
			unsigned short mode;
			unsigned short foot;
		};


		//加密
		//失败抛WinException异常
		//用户中断返回false
		template<int CryMode, class LamT>
		static void EncryptFile(df::CC fromFile, df::CC toFile, df::CC password, LamT && lam)
		{
			df::File fi, fiTo;

			if (!fi.Open(fromFile, false, false, true))
				Throw_df(fromFile + tcc_("\r\n文件打开失败!"));
			if (!fiTo.Open(toFile, true, true))
				Throw_df(toFile + tcc_("\r\n文件创建失败!"));

			std::string buf;
			buf.reserve(10 * 1024);
			CryptAlg<CryMode> crypt;
			//文件头
			CryptHeader head;
			head.mode = CryMode;
			head.foot = fi.GetFileSize() % crypt.KeySize;
			if (head.foot > 0)
			{
				head.foot = crypt.KeySize - head.foot;
			}

			if (!fiTo.Write(&head, sizeof(head)))
				ThrowLastErr(toFile + tcc_("\r\n文件写入失败!"));

			//初始化密钥
			crypt.InitKey(password);

			long long succedByte = 0;
			while (fi.Read(&buf[0], buf.capacity()))
			{
				short foot = fi.succeedByte_ % crypt.KeySize;
				if (foot > 0)
				{
					foot = crypt.KeySize - foot;
					head.foot = foot;
					memset(&buf[0] + fi.succeedByte_, 0, foot);
					fi.succeedByte_ += foot;
				}
				if (!crypt.Encrypt(&buf[0], &buf[0], fi.succeedByte_))
					Throw_df(tcc_("加密失败!"));

				if (!fiTo.Write(&buf[0], fi.succeedByte_))
					ThrowLastErr(toFile + tcc_("\r\n文件写入失败!"));

				succedByte += fi.succeedByte_;
				if (!df::ExecuteFunctor(lam, fi.GetFileSize(), succedByte))
					break;
			}
		}

		template <int CryMode>
		static void EncryptFile(df::CC fromFile, df::CC toFile, df::CC password)
		{
			EncryptFile <CryMode>(fromFile, toFile, password, [](long long, long long){});
		}


		//解密
		//失败抛WinException异常
		template<class LamT>
		static void DecryptFile(df::CC fromFile, df::CC toFile, df::CC password, LamT lam)
		{
			df::File fi;


			if (!fi.Open(fromFile, false, false, true))
				Throw_df(fromFile + tcc_("\r\n文件打开失败!"));

			CryptHeader head = { 0 };

			if (!fi.Read(&head, sizeof(head)))
				ThrowLastErr(fromFile + tcc_("\r\n文件头读取失败!"));

			if (fi.succeedByte_ < sizeof(head) || head.foot >= 16
				|| head.mode<CryptMode::DES_CBC || head.mode>CryptMode::AES_CBC)
			{
				Throw_df(fromFile + tcc_("\r\n此文件未加密!"));
			}


			df::File fiTo;

			if (!fiTo.Open(toFile, true, true))
				Throw_df(toFile + tcc_("\r\n文件创建失败!"));


			if (head.mode < CryptMode::DES_CBC || head.mode>CryptMode::AES_CBC)
				Throw_df(tcc_("未知算法"));


			typedef void(*DecryptFunc)(const df::CC &, df::File &, df::File &, const LamT &);
			DecryptFunc deFunc[] = {
				WriteDecrypt<CryptMode::DES_CBC, LamT>
				, WriteDecrypt<CryptMode::Triple_DES_CBC, LamT>
				, WriteDecrypt<CryptMode::AES_CBC, LamT>
			};

			deFunc[head.mode - CryptMode::DES_CBC](password, fi, fiTo, lam);


			fiTo.SetFileSize(fi.GetFileSize() - head.foot - sizeof(head));

		}

		static void DecryptFile(df::CC fromFile, df::CC toFile, df::CC password)
		{
			DecryptFile(fromFile, toFile, password, [](long long, long long){});
		}


	private:
		template<int CAlg, class LamT>
		static void WriteDecrypt(const df::CC & password, df::File & fi, df::File & fiTo, const LamT & lam)
		{
			std::string buf;
			buf.reserve(10 * 1024);
			CryptAlg<CAlg> crypt;
			crypt.InitKey(password);
			long long succedByte = 0;
			while (fi.Read(&buf[0], buf.capacity()))
			{
				if (!crypt.Decrypt(&buf[0], &buf[0], fi.succeedByte_))
					Throw_df(tcc_("解密失败!"));

				if (!fiTo.Write(&buf[0], fi.succeedByte_))
					ThrowLastErr(tcc_("文件写入失败!"));

				succedByte += fi.succeedByte_;
				if (!df::ExecuteFunctor(lam, fi.GetFileSize(), succedByte))
					break;
			}
		}

	};
}
